// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/socket_address.h>
#include <melon/init/init.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <stdlib.h>

#include <pollux/common/base/exceptions.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/checked_arithmetic.h>
#include <pollux/functions/prestosql/arithmetic.h>
#include <pollux/functions/prestosql/fail.h>
#include <pollux/functions/prestosql/string_functions.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/remote/client/remote.h>
#include <pollux/functions/remote/if/gen-cpp2/RemoteFunctionService.h>
#include <pollux/functions/remote/server/remote_function_service.h>
#include <pollux/functions/remote/utils/remote_function_service_provider.h>
#include <pollux/serializers/presto_serializer.h>

using ::apache::thrift::ThriftServer;
using ::kumo::pollux::test::assertEqualVectors;

namespace kumo::pollux::functions {
namespace {

struct Foo {
  explicit Foo(int64_t id) : id_(id) {}

  int64_t id() const {
    return id_;
  }

  int64_t id_;

  static std::string serialize(const std::shared_ptr<Foo>& foo) {
    return std::to_string(foo->id_);
  }

  static std::shared_ptr<Foo> deserialize(const std::string& serialized) {
    return std::make_shared<Foo>(std::stoi(serialized));
  }
};

template <typename T>
struct OpaqueTypeFunction {
  template <typename TInput, typename TOutput>
  MELON_ALWAYS_INLINE void call(TOutput& result, const TInput& a) {
    KLOG(INFO) << "OpaqueTypeFunction.value: " << a->id();
    result = a->id();
  }
};

// Parametrize in the serialization format so we can test both presto page and
// unsafe row.
class RemoteFunctionTest
    : public functions::test::FunctionBaseTest,
      public ::testing::WithParamInterface<remote::PageFormat> {
 public:
  void SetUp() override {
    auto params = startLocalThriftServiceAndGetParams();
    registerRemoteFunctions(params);
  }

  void TearDown() override {
    OpaqueType::clearSerializationRegistry();
  }

  // Registers a few remote functions to be used in this test.
  void registerRemoteFunctions(RemoteFunctionServiceParams params) {
    RemoteVectorFunctionMetadata metadata;
    metadata.serdeFormat = GetParam();
    metadata.location = params.serverAddress;

    // Register the remote adapter.
    auto plusSignatures = {exec::FunctionSignatureBuilder()
                               .returnType("bigint")
                               .argumentType("bigint")
                               .argumentType("bigint")
                               .build()};
    registerRemoteFunction("remote_plus", plusSignatures, metadata);

    auto failSignatures = {exec::FunctionSignatureBuilder()
                               .returnType("unknown")
                               .argumentType("integer")
                               .argumentType("varchar")
                               .build()};
    registerRemoteFunction("remote_fail", failSignatures, metadata);

    RemoteVectorFunctionMetadata wrongMetadata = metadata;
    wrongMetadata.location = melon::SocketAddress(); // empty address.
    registerRemoteFunction("remote_wrong_port", plusSignatures, wrongMetadata);

    auto divSignatures = {exec::FunctionSignatureBuilder()
                              .returnType("double")
                              .argumentType("double")
                              .argumentType("double")
                              .build()};
    registerRemoteFunction("remote_divide", divSignatures, metadata);

    auto substrSignatures = {exec::FunctionSignatureBuilder()
                                 .returnType("varchar")
                                 .argumentType("varchar")
                                 .argumentType("integer")
                                 .build()};
    registerRemoteFunction("remote_substr", substrSignatures, metadata);

    auto opaqueSignatures = {exec::FunctionSignatureBuilder()
                                 .returnType("bigint")
                                 .argumentType("opaque")
                                 .build()};
    registerRemoteFunction("remote_opaque", opaqueSignatures, metadata);

    // Registers the actual function under a different prefix. This is only
    // needed for tests since the thrift service runs in the same process.
    register_function<PlusFunction, int64_t, int64_t, int64_t>(
        {params.functionPrefix + ".remote_plus"});
    register_function<FailFunction, UnknownValue, int32_t, Varchar>(
        {params.functionPrefix + ".remote_fail"});
    register_function<CheckedDivideFunction, double, double, double>(
        {params.functionPrefix + ".remote_divide"});
    register_function<SubstrFunction, Varchar, Varchar, int32_t>(
        {params.functionPrefix + ".remote_substr"});
    register_function<OpaqueTypeFunction, int64_t, std::shared_ptr<Foo>>(
        {params.functionPrefix + ".remote_opaque"});

    registerOpaqueType<Foo>("Foo");
    OpaqueType::registerSerialization<Foo>(
        "Foo", Foo::serialize, Foo::deserialize);
  }
};

TEST_P(RemoteFunctionTest, simple) {
  auto inputVector = make_flat_vector<int64_t>({1, 2, 3, 4, 5});
  auto results = evaluate<SimpleVector<int64_t>>(
      "remote_plus(c0, c0)", make_row_vector({inputVector}));

  auto expected = make_flat_vector<int64_t>({2, 4, 6, 8, 10});
  assertEqualVectors(expected, results);
}

TEST_P(RemoteFunctionTest, string) {
  auto inputVector =
      make_flat_vector<StringView>({"hello", "my", "remote", "world"});
  auto inputVector1 = make_flat_vector<int32_t>({2, 1, 3, 5});
  auto results = evaluate<SimpleVector<StringView>>(
      "remote_substr(c0, c1)", make_row_vector({inputVector, inputVector1}));

  auto expected = make_flat_vector<StringView>({"ello", "my", "mote", "d"});
  assertEqualVectors(expected, results);
}

TEST_P(RemoteFunctionTest, tryException) {
  // remote_divide throws if denominator is 0.
  auto numeratorVector = make_flat_vector<double>({0, 1, 4, 9, 16});
  auto denominatorVector = make_flat_vector<double>({0, 1, 2, 3, 4});
  auto data = make_row_vector({numeratorVector, denominatorVector});
  auto results =
      evaluate<SimpleVector<double>>("TRY(remote_divide(c0, c1))", data);

  ASSERT_EQ(results->size(), 5);
  auto expected = make_flat_vector<double>({0 /* doesn't matter*/, 1, 2, 3, 4});
  expected->set_null(0, true);

  assertEqualVectors(expected, results);
}

TEST_P(RemoteFunctionTest, conditionalConjunction) {
  // conditional conjunction disables throwing on error.
  auto inputVector0 = make_flat_vector<bool>({true, true});
  auto inputVector1 = make_flat_vector<int32_t>({1, 2});
  auto data = make_row_vector({inputVector0, inputVector1});
  auto results = evaluate<SimpleVector<StringView>>(
      "case when (c0 OR remote_fail(c1, 'error')) then 'hello' else 'world' end",
      data);

  ASSERT_EQ(results->size(), 2);
  auto expected = make_flat_vector<StringView>({"hello", "hello"});
  assertEqualVectors(expected, results);
}

TEST_P(RemoteFunctionTest, tryErrorCode) {
  // remote_fail doesn't throw, but returns error code.
  auto errorCodesVector = make_flat_vector<int32_t>({1, 2});
  auto errorMessagesVector =
      make_flat_vector<StringView>({"failed 1", "failed 2"});
  auto data = make_row_vector({errorCodesVector, errorMessagesVector});
  exec::ExprSet exprSet(
      {makeTypedExpr("TRY(remote_fail(c0, c1))", as_row_type(data->type()))},
      &execCtx_);
  std::optional<SelectivityVector> rows;
  exec::EvalCtx context(&execCtx_, &exprSet, data.get());
  std::vector<VectorPtr> results(1);
  SelectivityVector defaultRows(data->size());
  exprSet.eval(defaultRows, context, results);

  ASSERT_EQ(results[0]->size(), 2);
}

TEST_P(RemoteFunctionTest, opaque) {
  // TODO: Support opaque type serialization in SPARK_UNSAFE_ROW
  if (GetParam() == remote::PageFormat::SPARK_UNSAFE_ROW) {
    GTEST_SKIP()
        << "opaque type serialization not supported in SPARK_UNSAFE_ROW";
  }
  auto inputVector = make_flat_vector<std::shared_ptr<void>>(
      2,
      [](vector_size_t row) { return std::make_shared<Foo>(row + 10); },
      /*is_null_at=*/nullptr,
      pollux::OPAQUE<Foo>());
  auto data = make_row_vector({inputVector});
  KLOG(INFO) << "type of data = " << data->type()->toString();

  auto results = evaluate<SimpleVector<int64_t>>(
      "remote_opaque(c0)", make_row_vector({inputVector}));

  auto expected = make_flat_vector<int64_t>({10, 11});
  assertEqualVectors(expected, results);
}

TEST_P(RemoteFunctionTest, connectionError) {
  auto inputVector = make_flat_vector<int64_t>({1, 2, 3, 4, 5});
  auto func = [&]() {
    evaluate<SimpleVector<int64_t>>(
        "remote_wrong_port(c0, c0)", make_row_vector({inputVector}));
  };

  // Check it throw and that the exception has the "connection refused"
  // substring.
  EXPECT_THROW(func(), PolluxRuntimeError);
  try {
    func();
  } catch (const PolluxRuntimeError& e) {
    EXPECT_THAT(e.message(), testing::HasSubstr("Channel is !good()"));
  }
}

POLLUX_INSTANTIATE_TEST_SUITE_P(
    RemoteFunctionTestFixture,
    RemoteFunctionTest,
    ::testing::Values(
        remote::PageFormat::PRESTO_PAGE,
        remote::PageFormat::SPARK_UNSAFE_ROW));

} // namespace
} // namespace kumo::pollux::functions

int main(int argc, char** argv) {
  testing::InitGoogleTest(&argc, argv);
  melon::Init init{&argc, &argv, false};
  return RUN_ALL_TESTS();
}
